rm(list=ls())
library(np)
library(xtable)
library(pastecs)


#################################
#  PROCEDURE "mediation.late.x" #
#################################
# INPUT
#y: outcome
#d: treatment (binary)
#m: mediator (continuous)
#zd: instrument for treatment (binary)
#zm: instrument for mediator (continuous)
#x: covariates
#trim: trimming level
#nonpara: if 1, first step estimators are estimated nonparametrically, otherwise parametrically 
#csquared: if nonpara is not 1 such that parametric first step estimators are used, setting csquared=1 means that both the control function c and its squared term are used as regressors, otherwise only c
#ruleofthumb: if equal to 1, the rule of thumb is used for bandwidth selection in nonparametric regression, otherwise least squares cross-validation
#boot: number of bootstrap replications for inference 
# bw1,...bw8: user provided bandwidths for nonparametric first step estimators
# minm: minimum number of observations when evaluating the conditional distribution function c

mediation.late.x<-function(y,d,m,zd, zm, x, trim=0.05, nonpara=1, csquared=0, ruleofthumb=1, boot=1999, bw1=NULL, bw2=NULL, bw3=NULL, bw4=NULL, bw5=NULL, bw6=NULL, bw7=NULL, bw8=NULL, minm=40){
  temp<-effects.late.x(y=y,d=d,m=m,zd=zd, zm=zm, x=x, trim=trim, nonpara=nonpara, csquared=csquared, ruleofthumb=ruleofthumb, bw1=bw1, bw2=bw2, bw3=bw3, bw4=bw4, bw5=bw5, bw6=bw6, bw7=bw7, bw8=bw8, minm=minm)
  if (is.null(bw1)) bw1<-temp$bw1
  if (is.null(bw2)) bw2<-temp$bw2
  if (is.null(bw3)) bw3<-temp$bw3
  if (is.null(bw4)) bw4<-temp$bw4
  if (is.null(bw5)) bw5<-temp$bw5
  if (is.null(bw6)) bw6<-temp$bw6
  if (is.null(bw7)) bw7<-temp$bw7
  if (is.null(bw8)) bw8<-temp$bw8
  temp2<-bootstrap.mediation.late.x(y=y,d=d,m=m,zd=zd,zm=zm, x=x, boot=boot,trim=trim, nonpara=nonpara, csquared=csquared, bw1=bw1, bw2=bw2, bw3=bw3, bw4=bw4, bw5=bw5, bw6=bw6, bw7=bw7, bw8=bw8, minm=minm)
  list(mc=temp2$mc, te=temp$te,  de.treat=temp$de.treat,  de.control=temp$de.control, ie.treat=temp$ie.treat, ie.control=temp$ie.control, de.treat.trim=temp$de.treat.trim, de.control.trim=temp$de.control.trim, ie.treat.trim=temp$ie.treat.trim, ie.control.trim=temp$ie.control.trim,   de.para=temp$de.para, ie.para=temp$ie.para,  te.trim=temp$te.trim, sd.te=temp2$sd.te,  sd.de.treat=temp2$sd.de.treat,  sd.de.control=temp2$sd.de.control,  sd.ie.treat=temp2$sd.ie.treat,  sd.ie.control=temp2$sd.ie.control, sd.de.treat.trim=temp2$sd.de.treat.trim, sd.de.control.trim=temp2$sd.de.control.trim,  sd.ie.treat.trim=temp2$sd.ie.treat.trim, sd.ie.control.trim=temp2$sd.ie.control.trim, sd.de.para=temp2$sd.de.para, sd.ie.para=temp2$sd.ie.para, sd.te.trim=temp2$sd.te.trim, bw1=bw1, bw2=bw2, bw3=bw3, bw4=bw4, bw5=bw5, bw6=bw6, bw7=bw7, bw8=bw8)
}


effects.late.x<-function(y,d,m,zd,  x, zm, trim=0.05, nonpara=1, csquared=0, ruleofthumb=1, bw1=NULL, bw2=NULL, bw3=NULL, bw4=NULL, bw5=NULL, bw6=NULL, bw7=NULL, bw8=NULL, minm=40){
dzd=d*zd
one_dzd=(1-d)*(1-zd)

if (nonpara!=1) {
pscore2=glm(zd~x,family=binomial(probit))$fitted
pscore2[pscore2>0.9999]=0.9999
pscore2[pscore2<0.0001]=0.0001
}
if (nonpara==1) {
if (is.null(bw1)) {
if (ruleofthumb!=1) bw1<-npregbw(ydat=factor(zd), xdat=x, regtype="lc", ckertype="gaussian", bwmethod="cv.ls")$bw
if (ruleofthumb==1) bw1<-npudensbw(dat=x, ckertype="gaussian", bwmethod="normal-reference")$bw
}
pscore2<-(npreg(bws=bw1, tydat=factor(zd), txdat=x, regtype="lc", ckertype="gaussian")$mean)-1
pscore2[pscore2>0.9999]=0.9999
pscore2[pscore2<0.0001]=0.0001
}

c.d=d*(zd-pscore2)
c.d_1=(d-1)*(zd-pscore2)
if (nonpara==1) {
zmx<-data.frame(c(zm),c(x))
x<-as.matrix(x)
zm<-as.matrix(zm)
if (is.null(bw2) | is.null(bw3)) {
if (ruleofthumb!=1) temp<-npcdensbw(ydat=m, xdat=zmx, ckertype="gaussian", bwmethod="cv.ls")
if (ruleofthumb==1) temp<-npcdensbw(ydat=m, xdat=zmx, ckertype="gaussian", bwmethod="normal-reference")
}
if (is.null(bw2))  bw2<-temp$xbw
if (is.null(bw3))  bw3<-temp$ybw
m.dist=npcdist(bws=c(bw3,bw2), tydat=m, txdat=zmx, regtype="lc", ckertype="gaussian")$condist
mm<-sort(m)

if (is.null(bw4)) {
if (ruleofthumb!=1) bw4<-npregbw(ydat=c.d, xdat=zmx, regtype="ll", ckertype="gaussian", bwmethod="cv.ls")$bw
if (ruleofthumb==1) bw4<-npudensbw(dat=zmx, ckertype="gaussian", bwmethod="normal-reference")$bw
}
c.den.d=npreg(bws=bw4, tydat=c.d, txdat=zmx, regtype="ll", ckertype="gaussian")$mean
if (is.null(bw5)) {
if (ruleofthumb!=1) bw5<-npregbw(ydat=c.d_1, xdat=zmx, regtype="ll", ckertype="gaussian", bwmethod="cv.ls")$bw
if (ruleofthumb==1) bw5<-bw4
}
c.num<-c()
for (i in 1:nrow(zmx)){
ind= ((m<=m[i]) | (m<=mm[minm]))
n<-length(m); nn<-sum(ind);
c.num.d=npreg(bws=bw4*(nn^((-1)/(4+ncol(zmx))))/(n^((-1)/(4+ncol(zmx)))), tydat=c.d[ind==1], txdat=zmx[ind==1,], eydat=c.d[i], exdat=zmx[i,], regtype="ll", ckertype="gaussian")$mean
if (is.na(c.num.d[i])) c.num.d=npreg(bws=2*bw4*(nn^((-1)/(4+ncol(zmx))))/(n^((-1)/(4+ncol(zmx)))), tydat=c.d[ind==1], txdat=zmx[ind==1,], eydat=c.d[i], exdat=zmx[i,],   regtype="ll", ckertype="gaussian")$mean
if (is.na(c.num.d[i])) c.num.d=npreg(bws=4*bw4*(nn^((-1)/(4+ncol(zmx))))/(n^((-1)/(4+ncol(zmx)))), tydat=c.d[ind==1], txdat=zmx[ind==1,], eydat=c.d[i], exdat=zmx[i,],  regtype="ll", ckertype="gaussian")$mean
if (is.na(c.num.d[i])) c.num.d=npreg(bws=8*bw4*(nn^((-1)/(4+ncol(zmx))))/(n^((-1)/(4+ncol(zmx)))), tydat=c.d[ind==1], txdat=zmx[ind==1,], eydat=c.d[i], exdat=zmx[i,],  regtype="ll", ckertype="gaussian")$mean
if (is.na(c.num.d[i])) c.num.d=c(1,zm[i,],x[i,])%*%(lm(c.d[ind==1]~zm[ind==1,]+x[ind==1,])$coef)
c.num.d_1=npreg(bws=bw5*(nn^((-1)/(4+ncol(zmx))))/(n^((-1)/(4+ncol(zmx)))), tydat=c.d_1[ind==1], txdat=zmx[ind==1,],  eydat=c.d_1[i], exdat=zmx[i,],   regtype="ll", ckertype="gaussian")$mean
if (is.na(c.num.d_1[i])) c.num.d_1=npreg(bws=2*bw5*(nn^((-1)/(4+ncol(zmx))))/(n^((-1)/(4+ncol(zmx)))), tydat=c.d_1[ind==1], txdat=zmx[ind==1,],  eydat=c.d_1[i], exdat=zmx[i,],   regtype="ll", ckertype="gaussian")$mean
if (is.na(c.num.d_1[i])) c.num.d_1=npreg(bws=4*bw5*(nn^((-1)/(4+ncol(zmx))))/(n^((-1)/(4+ncol(zmx)))), tydat=c.d_1[ind==1], txdat=zmx[ind==1,],  eydat=c.d_1[i], exdat=zmx[i,],   regtype="ll", ckertype="gaussian")$mean
if (is.na(c.num.d_1[i])) c.num.d_1=npreg(bws=8*bw5*(nn^((-1)/(4+ncol(zmx))))/(n^((-1)/(4+ncol(zmx)))), tydat=c.d_1[ind==1], txdat=zmx[ind==1,], eydat=c.d_1[i], exdat=zmx[i,],   regtype="ll", ckertype="gaussian")$mean
if (is.na(c.num.d_1[i])) c.num.d_1=c(1,zm[i,],x[i,])%*%(lm(c.d_1[ind==1]~zm[ind==1,]+x[ind==1,])$coef)
c.num<-c(c.num, (d[i]*c.num.d+(1-d[i])*c.num.d_1))
}
}
#
if (nonpara!=1){
c.den.d=lm(c.d~zm+x)$fitted
if (is.null(bw2) | is.null(bw3)) temp<-npcdensbw(ydat=m, xdat=cbind(zm,x), ckertype="gaussian", bwmethod="normal-reference")
if (is.null(bw2))  bw2<-temp$xbw
if (is.null(bw3))  bw3<-temp$ybw
m.dist=npcdist(bws=c(bw3,bw2), tydat=m, txdat=cbind(zm,x), ckertype="gaussian")$condist
mm<-sort(m)
x<-as.matrix(x)
zm<-as.matrix(zm)
c.num<-c()
for (i in 1:nrow(x)){
ind= ((m<=m[i]) | (m<=mm[minm]))
c.num.d=lm(c.d[ind==1]~zm[ind==1,]+x[ind==1,])
c.num.d_1=lm(c.d_1[ind==1]~zm[ind==1,]+x[ind==1,])
c.num<-c(c.num, (d[i]*(c(1,zm[i,],x[i,])%*%c.num.d$coef)+(1-d[i])*(c(1,zm[i,],x[i,])%*%c.num.d_1$coef)))
}
}
c=c.num/c.den.d*m.dist

if (nonpara!=1){
if (csquared!=1 ){
pscore1<-glm(zd~cbind(c,m,x),family=binomial(probit))$fitted.values
pscored<-glm(d~cbind(c,m,x),family=binomial(probit))$fitted.values
pscored[pscored<0]=0; pscored[pscored>1]=1
pscoredz<-glm(dzd~cbind(c,m,x),family=binomial(probit))$fitted.values
pscoreone_dz<-glm(one_dzd~cbind(c,m,x),family=binomial(probit))$fitted.values
pscoreone_dz[pscoreone_dz<0]=0; pscoreone_dz[pscoreone_dz>1]=1
}
if (csquared==1){
c2<-c^2
cx<-c*x
pscore1<-glm(zd~cbind(c,c2,m, m^2, x),family=binomial(probit))$fitted.values
pscored<-glm(d~cbind(c, c2,m, m^2, x),family=binomial(probit))$fitted.values
pscored[pscored<0]=0; pscored[pscored>1]=1
pscoredz<-glm(dzd~cbind(c,c2,m,m^2, x),family=binomial(probit))$fitted.values
pscoreone_dz<-glm(one_dzd~cbind(c,c2,m,m^2, x),family=binomial(probit))$fitted.values
pscoreone_dz[pscoreone_dz<0]=0; pscoreone_dz[pscoreone_dz>1]=1
}
}
if (nonpara==1){
cmx<-data.frame(c(c),c(m),c(x))

if (is.null(bw6)) {
if (ruleofthumb!=1) bw6<-npregbw(ydat=factor(zd), xdat=cmx, regtype="lc", ckertype="gaussian", bwmethod="cv.ls")$bw
if (ruleofthumb==1) bw6<-npudensbw(dat=cmx, ckertype="gaussian", bwmethod="normal-reference")$bw
}
pscore1<-(npreg(bws=bw6, tydat=factor(zd), txdat=cmx, regtype="lc", ckertype="gaussian")$mean)-1

if (is.null(bw7)) {
if (ruleofthumb!=1) bw7<-npregbw(ydat=factor(d), xdat=cmx, regtype="lc", ckertype="gaussian", bwmethod="cv.ls")$bw
if (ruleofthumb==1) bw7<-bw6
}
pscored<-(npreg(bws=bw7, tydat=factor(d), txdat=cmx, regtype="lc", ckertype="gaussian")$mean)-1
if (is.null(bw8)) {
if (ruleofthumb!=1) bw8<-npregbw(ydat=factor(dzd), xdat=cmx, regtype="lc", ckertype="gaussian", bwmethod="cv.ls")$bw
if (ruleofthumb==1) bw8<-bw7
}
pscoredz<-(npreg(bws=bw8, tydat=factor(dzd), txdat=cmx, regtype="lc", ckertype="gaussian")$mean)-1
}

firststage=mean(d/pscore2*(zd-pscore2)/(1-pscore2))
omega=1-(pscore1-pscore2)/(pscoredz-pscored*pscore2)
one_omega=1/omega
omega[omega<0.01]=0.01
omega[omega>100]=100
one_omega[one_omega<0.01]=0.01
one_omega[one_omega>100]=100
wgt<-(zd-pscore2)/(pscore2*(1-pscore2))
y1m0=mean(y*d*omega*wgt)/firststage
y1m1=mean( y*d*wgt )/firststage
y0m1=mean(y*(d-1)*one_omega*wgt )/firststage
y0m0= mean(y*(d-1)*wgt)/firststage
ie.treat=y1m1-y1m0
ie.control=y0m1-y0m0
de.treat=y1m1-y0m1
de.control=y1m0-y0m0
te=y1m1-y0m0

pred.d<-fitted.values(glm(d~zd+x,family=binomial(probit)))
pred.m<-fitted.values(lm(m~zm+pred.d+x))
temp<-lm(y~cbind(pred.d,pred.m,x))
temp2<-lm(m~cbind(pred.d,x))$coef[2]

ind=1-(d*omega*wgt/sum(d*omega*wgt)>trim) 
y1m0.trim=mean(y[ind==1]*d[ind==1]*omega[ind==1]*wgt[ind==1])/firststage
ind= 1-(d*wgt/sum(d*wgt)>trim)
y1m1.trim=mean( y[ind==1]*d[ind==1]*wgt[ind==1] )/firststage
ind=1-((d-1)*one_omega*wgt/sum((d-1)*one_omega*wgt)>trim)
y0m1.trim=mean(y[ind==1]*(d[ind==1]-1)*one_omega[ind==1]*wgt[ind==1] )/firststage
ind=1-((d-1)*wgt/sum((d-1)*wgt)>trim)
y0m0.trim= mean(y[ind==1]*(d[ind==1]-1)*wgt[ind==1])/firststage
ie.treat.trim=y1m1.trim-y1m0.trim
ie.control.trim=y0m1.trim-y0m0.trim
de.treat.trim=y1m1.trim-y0m1.trim
de.control.trim=y1m0.trim-y0m0.trim
te.trim=y1m1.trim-y0m0.trim
list(te=te, de.treat=de.treat, de.control=de.control,  ie.treat=ie.treat, ie.control=ie.control, de.treat.trim=de.treat.trim, de.control.trim=de.control.trim, ie.treat.trim=ie.treat.trim, ie.control.trim=ie.control.trim, de.para=temp$coef[2], ie.para=temp$coef[3]*temp2, te.trim=te.trim, bw1=bw1, bw2=bw2, bw3=bw3, bw4=bw4, bw5=bw5, bw6=bw6, bw7=bw7, bw8=bw8,c=c)
}


bootstrap.mediation.late.x<-function(y,d,m,zd,zm,x, boot=1999,trim=0.05, nonpara=1, csquared=0, ruleofthumb=1, bw1=NULL, bw2=NULL, bw3=NULL, bw4=NULL, bw5=NULL, bw6=NULL, bw7=NULL, bw8=NULL, minm=40){
obs<-length(y)
mc=c()
temp=c()
while(length(temp)<boot){
sboot<-sample(1:obs,obs,TRUE)
yb=y[sboot]
db<-d[sboot]
zdb<-zd[sboot]
mb=m[sboot]
if (length(zm)==length(y)) zmb<-zm[sboot]
if (length(zm)!=length(y)) zmb<-zm[sboot,]
if (length(x)==length(y)) xb<-x[sboot]
if (length(x)!=length(y)) xb<-x[sboot,]
est<-c(effects.late.x(yb,db,mb,zdb, zmb, xb, trim=trim, nonpara=nonpara, csquared=csquared, ruleofthumb=ruleofthumb, bw1=bw1, bw2=bw2, bw3=bw3, bw4=bw4, bw5=bw5, bw6=bw6, bw7=bw7, bw8=bw8,minm=minm))
if (sum(is.na(est))==0) mc<-rbind(mc, est)
temp<-c(temp,1)
}
list(mc=mc, sd.te=sd(as.numeric(mc[,1])), sd.de.treat=sd(as.numeric(mc[,2])), sd.de.control=sd(as.numeric(mc[,3])), sd.ie.treat=sd(as.numeric(mc[,4])), sd.ie.control=sd(as.numeric(mc[,5])), sd.de.treat.trim=sd(as.numeric(mc[,6])), sd.de.control.trim=sd(as.numeric(mc[,7])),  sd.ie.treat.trim=sd(as.numeric(mc[,8])),  sd.ie.control.trim=sd(as.numeric(mc[,9])), sd.de.para=sd(as.numeric(mc[,10])), sd.ie.para=sd(as.numeric(mc[,11]))   )
}

################
# APPLICATION  #
################

#load data set
load("C:/BHPS.RData")

# descriptive stats
descrvar=cbind(x,y,zd,zm)
xtable(  cbind( t(stat.desc(descrvar[d==1,])[9,]),t(stat.desc(descrvar[d==1,])[13,]), t(stat.desc(descrvar[d==0,])[9,]),t(stat.desc(descrvar[d==0,])[13,]),    t(stat.desc(descrvar[m>12000,])[9,]),t(stat.desc(descrvar[m>12000,])[13,]), t(stat.desc(descrvar[m<=12000,])[9,]),t(stat.desc(descrvar[m<=12000,])[13,]))  ,digits=3)

#semi-parametric estimation
est<-mediation.late.x(y=y,d=d,m=m,zd=zd, zm=zm, x=x, trim=0.05, boot=999, nonpara=0, csquared=0, ruleofthumb=1)
pval<-c()
for (i in 1:11) pval<-c(pval, c(2*(min( mean((as.numeric(est$mc[,i])-as.numeric(est[i+1]))<=as.numeric(est[i+1])) , mean((as.numeric(est$mc[,i])-as.numeric(est[i+1]))>as.numeric(est[i+1])) ) ) ) )
results<-rbind(cbind(est$te, est$de.treat, est$de.control, est$ie.treat,  est$ie.control, est$de.treat.trim, est$de.control.trim, est$ie.treat.trim,  est$ie.control.trim, est$de.para, est$ie.para), cbind(est$sd.te, est$sd.de.treat, est$sd.de.control, est$sd.ie.treat,  est$sd.ie.control, est$sd.de.treat.trim, est$sd.de.control.trim, est$sd.ie.treat.trim,  est$sd.ie.control.trim, est$sd.de.para, est$sd.ie.para), cbind(2*pnorm(-abs(est$te/est$sd.te)), 2*pnorm(-abs(est$de.treat/est$sd.de.treat)), 2*pnorm(-abs(est$de.control/est$sd.de.control)),  2*pnorm(-abs(est$ie.treat/est$sd.ie.treat)), 2*pnorm(-abs(est$ie.control/est$sd.ie.control)), 2*pnorm(-abs(est$de.treat.trim/est$sd.de.treat.trim)), 2*pnorm(-abs(est$de.control.trim/est$sd.de.control.trim)),  2*pnorm(-abs(est$ie.treat.trim/est$sd.ie.treat.trim)), 2*pnorm(-abs(est$ie.control.trim/est$sd.ie.control.trim)) ,2*pnorm(-abs(est$de.para/est$sd.de.para))   , 2*pnorm(-abs(est$ie.para/est$sd.ie.para))   ), pval  )
rownames(results)=c("effect","boot se","asymp. p-val", "boot p-val")
colnames(results)=c("late", "de.treat", "de.control",  "ie.treat", "ie.control", "de.treat.trim", "de.control.trim", "ie.treat.trim", "ie.control.trim", "de.para", "ie.para")

#results
xtable(results, digits=3)

